#!/usr/bin python3
# -*- coding:utf-8 -*-
import os
import numpy as np
import argparse

import torch
from torch.utils.data import DataLoader
from torchvision import models
from torch.autograd import Variable
import torchvision.transforms as transforms
from YOLODataset import YOLODataset
# from visualize import Visualizer

from backbones import resnet
from backbones import vggnet
from loss import YOLOLoss

import warnings
warnings.filterwarnings("ignore", category=UserWarning)


def make_parser():
    parser = argparse.ArgumentParser("YOLO-v1 Trainer Params Setting!")
    parser.add_argument('--train_imgs_root', type=str, default='', help='train images root directory')
    parser.add_argument('--test_imgs_root', type=str, default='', help='train images root directory')
    parser.add_argument('--use_gpu', type=bool, default=True, help='use cpu to train yolo-v1')
    parser.add_argument('--gpu_ids', type=str, default='0', help='set gpu id for training')
    parser.add_argument('--lr_rate', type=float, default=0.001, help='initial learning rate for optimizer')
    parser.add_argument('--num_epochs', type=int, default=50, help='train epochs')
    parser.add_argument('--batch_size', type=int, default=48, help='batch_size for one train')
    parser.add_argument('--use_resnet', type=bool, default=True, help='use resnet-50 as yolo-v1\'s backbone')
    args = parser.parse_args()
    return args

os.environ["CUDA_VISIBLE_DEVICES"] ="0"

def train():
    args = make_parser()

    # det_net = None
    if args.use_resnet:
        det_net = resnet.resnet50()
    else:
        det_net =vggnet.vgg16_bn()

    print("load pretrained model.")
    if args.use_resnet:
        resnet50 = models.resnet50(pretrained=True)
        resnet50_state_dict = resnet50.state_dict()
        det_net_state_dict = det_net.state_dict()
        for k in resnet50_state_dict.keys():
            if k in det_net_state_dict.keys() and not k.startswith('fc'):
                det_net_state_dict[k] = resnet50_state_dict[k]
        det_net.load_state_dict(det_net_state_dict)
    else:
        vgg = models.vgg16_bn(pretrained=True)
        vgg_state_dict = vgg.state_dict()
        det_net_state_dict = det_net.state_dict()
        for k in vgg_state_dict.keys():
            if k in det_net_state_dict.keys() and k.startswith('features'):
                det_net_state_dict[k] = vgg_state_dict[k]
        det_net.load_state_dict(det_net_state_dict)

    if False:
        det_net.load_state_dict(torch.load('best.pth'))
    print('cuda', torch.cuda.current_device(), torch.cuda.device_count())

    criterion = YOLOLoss(7, 2, 20, 5, 0.5)

    if args.use_gpu:
        det_net.cuda()

    det_net.train()

    params = []
    params_dict = dict(det_net.named_parameters())
    for key, val in params_dict.items():
        if key.startswith('feature'):
            params += [{'params': [val], 'lr':args.lr_rate*1}]
        else:
            params += [{'params': [val], 'lr':args.lr_rate}]

    optimizer = torch.optim.SGD(params, lr=args.lr_rate, momentum=0.9, weight_decay=5e-4)
    train_dataset = YOLODataset(root=args.train_imgs_root, list_file=['voc2012.txt'], is_train=True,
                                transform=[transforms.ToTensor()])
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=4)

    test_dataset = YOLODataset(root=args.test_imgs_root, list_file='voc2007test.txt', is_train=False, transform=[transforms.ToTensor()])
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
    print('The dataset has %d images' % (len(train_dataset)))
    print('The batch_size is %d' % args.batch_size)
    logfile = open('log.txt', 'w')

    num_iter = 0
    # vis = Visualizer(env='lp_yolov1')
    test_loss = np.inf

    for epoch in range(args.num_epochs):
        det_net.train()
        if epoch == 30:
            args.lr_rate = 0.0001
        if epoch == 40:
            args.lr_rate = 0.00001
        for param_group in optimizer.param_groups:
            param_group['lr'] = args.lr_rate

        total_loss = 0
        for i, (imgs, labels) in enumerate(train_loader):
            imgs = Variable(imgs)
            labels = Variable(labels)

            if args.use_gpu:
                imgs, labels = imgs.cuda(), labels.cuda()

            pred = det_net(imgs)
            # print(pred.size(), labels.size())
            loss = criterion(pred, labels)
            total_loss += loss.data

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 5 == 0:
                print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f' %
                      (epoch + 1, args.num_epochs, i+1, len(train_loader), loss.data, total_loss/(i+1)))
                num_iter += 1
                # vis.plot_train_val(loss_train=total_loss / (i+1))

        # Validation
        validation_loss = 0.0
        det_net.eval()
        for i, (imgs, labels) in enumerate(test_loader):
            imgs = Variable(imgs, volatile=True)
            labels = Variable(labels, volatile=True)
            if args.use_gpu:
                imgs, labels = imgs.cuda(), labels.cuda()

            pred = det_net(imgs)
            loss = criterion(pred, labels)
            validation_loss += loss.data
        validation_loss /= len(test_loader)
        # vis.plot_train_val(loss_val=validation_loss)

        if test_loss > validation_loss:
            test_loss = validation_loss
            print("BEST TEST LOSS %.5f" % test_loss)
            torch.save(det_net.state_dict(), 'best_0712.pth')
        logfile.writelines(str(epoch) + '\t' + str(validation_loss) + '\n')
        logfile.flush()
        torch.save(det_net.state_dict(), 'yolov1_0712.pth')


if __name__ == '__main__':
    print("========== START TRAIN YOLO-V1 ===========")
    train()
    print("=========== END TRAIN YOLO-V1 ============")
